#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：scheduling_suite.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/20 20:23
能不能帮我写一些关于车间调度和算法的代码吗，比如GA解决柔性作业车间调度问题，解决流水车间调度问题，解决混合流水车间调度问题，或者再找几个其他算法一起
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
scheduling_suite.py

包含：
1) 问题建模：
   - Flow Shop Problem (FSP)
   - Hybrid Flow Shop (HFS，多阶段并行机)
   - Flexible Job Shop (FJSP，多机可行集)
2) 算法：
   - GA_FSP（排列编码）、SA_FSP（模拟退火）
   - GA_HFS（排列+分配编码）
   - GA_FJSP（作业序列+机器分配编码，基于“重复作业id序列”的常用解码）
3) 通用工具：
   - 评估/甘特图/收敛曲线

依赖：numpy, matplotlib
"""

import math
import random
import numpy as np
import matplotlib.pyplot as plt
from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional

# ============ 通用工具 ============

def seed_all(seed=0):
    random.seed(seed); np.random.seed(seed)

def plot_convergence(best_hist, title="Convergence"):
    plt.figure(figsize=(6,4))
    plt.plot(best_hist, lw=2)
    plt.xlabel("Generation / Iter"); plt.ylabel("Best makespan")
    plt.title(title); plt.grid(True, alpha=0.3); plt.tight_layout(); plt.show()

def plot_gantt(intervals_by_machine: Dict[str, List[Tuple[float,float,str]]], title="Gantt"):
    # intervals_by_machine: name -> list of (start, duration, label)
    machines = list(intervals_by_machine.keys())
    fig, ax = plt.subplots(figsize=(12, 0.7*len(machines)+3))
    yticks=[]; ylabels=[]
    for idx, m in enumerate(machines):
        items = sorted(intervals_by_machine[m], key=lambda x: x[0])
        for (st, du, lab) in items:
            ax.barh(idx, du, left=st)
            ax.text(st+du/2, idx, lab, ha="center", va="center", fontsize=8)
        yticks.append(idx); ylabels.append(m)
    ax.set_yticks(yticks); ax.set_yticklabels(ylabels)
    ax.set_xlabel("Time"); ax.set_title(title)
    ax.grid(axis="x", alpha=0.3)
    plt.tight_layout(); plt.show()

# ============ FSP（流水车间：所有作业相同机器顺序） ============

@dataclass
class FSP:
    # p[j, m]: 作业 j 在机器 m 上的加工时间
    p: np.ndarray   # shape [n_jobs, n_machines]

    @property
    def n_jobs(self): return self.p.shape[0]
    @property
    def n_machines(self): return self.p.shape[1]

    def makespan(self, perm: List[int]) -> float:
        """标准动态规划计算 Cmax"""
        n, m = self.n_jobs, self.n_machines
        perm = list(perm)
        C = np.zeros((len(perm), m), dtype=float)
        for i, j in enumerate(perm):
            for k in range(m):
                up = C[i-1, k] if i>0 else 0.0
                left = C[i, k-1] if k>0 else 0.0
                C[i, k] = max(up, left) + self.p[j, k]
        return float(C[-1, -1])

    def decode_gantt(self, perm: List[int]) -> Dict[str, List[Tuple[float,float,str]]]:
        """用于可视化甘特图"""
        n, m = self.n_jobs, self.n_machines
        C = np.zeros((len(perm), m), dtype=float)
        S = np.zeros((len(perm), m), dtype=float)
        for i, j in enumerate(perm):
            for k in range(m):
                upC = C[i-1, k] if i>0 else 0.0
                leftC = C[i, k-1] if k>0 else 0.0
                start = max(upC, leftC)
                dur = self.p[j, k]
                S[i, k] = start
                C[i, k] = start + dur
        gantt = {}
        for k in range(m):
            key = f"M{k}"
            gantt[key] = []
            for i, j in enumerate(perm):
                st = S[i, k]; du = self.p[j, k]
                gantt[key].append((st, du, f"J{j}"))
        return gantt

# ---- GA for FSP ----

class GA_FSP:
    def __init__(self, problem:FSP, pop=60, gens=300, pc=0.9, pm=0.2, seed=0):
        self.pb = problem; self.pop=pop; self.gens=gens
        self.pc=pc; self.pm=pm; seed_all(seed)

    def init_pop(self):
        pop=[]
        base = list(range(self.pb.n_jobs))
        for _ in range(self.pop):
            perm = base[:]; random.shuffle(perm); pop.append(perm)
        return pop

    def crossover_ox(self, a, b):
        """Order Crossover (OX)"""
        n=len(a)
        i, j = sorted(random.sample(range(n), 2))
        child = [None]*n
        child[i:j+1] = a[i:j+1]
        pos = (j+1)%n
        for k in range(n):
            x = b[(j+1+k)%n]
            if x not in child:
                child[pos] = x
                pos = (pos+1)%n
        return child

    def mutate_swap(self, a):
        n=len(a); i,j = random.sample(range(n),2)
        a[i],a[j] = a[j],a[i]

    def run(self, verbose=True):
        pop = self.init_pop()
        fitness = [self.pb.makespan(ind) for ind in pop]
        best_hist=[]
        for g in range(1, self.gens+1):
            # 选择（锦标赛）
            new=[]
            while len(new) < self.pop:
                i,j = random.randrange(self.pop), random.randrange(self.pop)
                p1 = pop[i] if fitness[i]<fitness[j] else pop[j]
                i,j = random.randrange(self.pop), random.randrange(self.pop)
                p2 = pop[i] if fitness[i]<fitness[j] else pop[j]
                c1, c2 = p1[:], p2[:]
                if random.random()<self.pc:
                    c1 = self.crossover_ox(p1,p2)
                    c2 = self.crossover_ox(p2,p1)
                if random.random()<self.pm: self.mutate_swap(c1)
                if random.random()<self.pm: self.mutate_swap(c2)
                new += [c1,c2]
            pop = new[:self.pop]
            fitness = [self.pb.makespan(ind) for ind in pop]
            best = min(fitness); best_hist.append(best)
            if verbose and g%50==0:
                print(f"[GA_FSP] Gen {g:4d} Best Cmax={best:.2f}")
        idx = int(np.argmin(fitness))
        return pop[idx], fitness[idx], best_hist

# ---- SA for FSP ----

class SA_FSP:
    def __init__(self, problem:FSP, T0=100.0, alpha=0.98, iters=5000, seed=0):
        self.pb=problem; self.T0=T0; self.alpha=alpha; self.iters=iters; seed_all(seed)

    def neighbor(self, perm):
        a = perm[:]
        i,j = random.sample(range(len(a)),2)
        a[i],a[j] = a[j],a[i]
        return a

    def run(self, verbose=True):
        cur = list(range(self.pb.n_jobs)); random.shuffle(cur)
        cur_val = self.pb.makespan(cur)
        best, best_val = cur[:], cur_val
        T=self.T0
        hist=[]
        for t in range(1, self.iters+1):
            nxt = self.neighbor(cur)
            nxt_val = self.pb.makespan(nxt)
            if (nxt_val < cur_val) or (random.random()<math.exp(-(nxt_val-cur_val)/max(1e-9,T))):
                cur, cur_val = nxt, nxt_val
                if nxt_val < best_val:
                    best, best_val = nxt[:], nxt_val
            hist.append(best_val)
            T *= self.alpha
            if verbose and t%1000==0:
                print(f"[SA_FSP] Iter {t:4d} Best Cmax={best_val:.2f}")
        return best, best_val, hist

# ============ HFS（混合流水车间：分阶段并行机） ============

@dataclass
class HFS:
    # stages: S 阶段；每阶段有 Ms[s] 台并行机器
    # proc_time[s][j]: 作业 j 在阶段 s 的加工时间（相同于所有并机）
    proc_time: List[np.ndarray]  # list of [n_jobs] arrays, len = S
    Ms: List[int]                # 每阶段并机数量

    @property
    def S(self): return len(self.proc_time)
    @property
    def n_jobs(self): return self.proc_time[0].shape[0]

    def decode(self, perm: List[int], assign: List[List[int]]) -> Tuple[float, Dict[str, List[Tuple[float,float,str]]]]:
        """
        assign[s][j] ∈ [0, Ms[s)-1] 表示作业 j 在阶段 s 选择的并机编号
        排程：每阶段按 perm 顺序，把作业分配到对应机器，机器就绪与作业就绪取 max
        """
        job_ready = np.zeros(self.n_jobs, dtype=float)
        mach_ready = [np.zeros(self.Ms[s], dtype=float) for s in range(self.S)]
        gantt = {f"S{s}M{m}":[] for s in range(self.S) for m in range(self.Ms[s])}
        for j in perm:
            for s in range(self.S):
                m = assign[s][j]
                st = max(job_ready[j], mach_ready[s][m])
                pt = float(self.proc_time[s][j])
                en = st + pt
                job_ready[j] = en; mach_ready[s][m] = en
                gantt[f"S{s}M{m}"].append((st, pt, f"J{j}"))
        Cmax = float(np.max(job_ready))
        return Cmax, gantt

# ---- GA for HFS ----

class GA_HFS:
    def __init__(self, problem:HFS, pop=60, gens=300, pc=0.9, pm=0.2, seed=0):
        self.pb=problem; self.pop=pop; self.gens=gens; self.pc=pc; self.pm=pm; seed_all(seed)

    def init_ind(self):
        perm = list(range(self.pb.n_jobs)); random.shuffle(perm)
        assign = [ [random.randrange(self.pb.Ms[s]) for _ in range(self.pb.n_jobs)] for s in range(self.pb.S) ]
        return (perm, assign)

    def crossover_perm_ox(self, a, b):
        n=len(a); i,j = sorted(random.sample(range(n),2))
        child = [None]*n; child[i:j+1]=a[i:j+1]
        pos=(j+1)%n
        for k in range(n):
            x=b[(j+1+k)%n]
            if x not in child:
                child[pos]=x; pos=(pos+1)%n
        return child

    def crossover_assign(self, A,B):
        # 均匀交叉每阶段每作业的机器分配
        C=[]
        for s in range(self.pb.S):
            row=[]
            for j in range(self.pb.n_jobs):
                row.append(A[s][j] if random.random()<0.5 else B[s][j])
            C.append(row)
        return C

    def mutate(self, perm, assign):
        # 交换两作业
        if random.random()<0.5:
            i,j = random.sample(range(len(perm)),2)
            perm[i],perm[j]=perm[j],perm[i]
        # 随机调整若干机器分配
        for s in range(self.pb.S):
            if random.random()<0.3:
                j = random.randrange(self.pb.n_jobs)
                assign[s][j] = random.randrange(self.pb.Ms[s])

    def run(self, verbose=True):
        pop = [self.init_ind() for _ in range(self.pop)]
        fitness = [self.pb.decode(ind[0], ind[1])[0] for ind in pop]
        best_hist=[]
        for g in range(1, self.gens+1):
            new=[]
            while len(new) < self.pop:
                i,j = random.randrange(self.pop), random.randrange(self.pop)
                p1 = pop[i] if fitness[i]<fitness[j] else pop[j]
                i,j = random.randrange(self.pop), random.randrange(self.pop)
                p2 = pop[i] if fitness[i]<fitness[j] else pop[j]
                c1=(p1[0][:], [r[:] for r in p1[1]])
                c2=(p2[0][:], [r[:] for r in p2[1]])
                if random.random()<self.pc:
                    c1 = (self.crossover_perm_ox(p1[0], p2[0]),
                          self.crossover_assign(p1[1], p2[1]))
                    c2 = (self.crossover_perm_ox(p2[0], p1[0]),
                          self.crossover_assign(p2[1], p1[1]))
                if random.random()<self.pm: self.mutate(*c1)
                if random.random()<self.pm: self.mutate(*c2)
                new += [c1,c2]
            pop = new[:self.pop]
            fitness = [self.pb.decode(ind[0], ind[1])[0] for ind in pop]
            best = min(fitness); best_hist.append(best)
            if verbose and g%50==0:
                print(f"[GA_HFS] Gen {g:4d} Best Cmax={best:.2f}")
        idx = int(np.argmin(fitness))
        return pop[idx], fitness[idx], best_hist

# ============ FJSP（柔性作业车间：工序可选多机） ============

@dataclass
class Operation:
    machine_time: Dict[int, int]  # 可行机 -> pt

@dataclass
class JobF:
    ops: List[Operation]

@dataclass
class FJSP:
    jobs: List[JobF]
    m: int

    @property
    def n_jobs(self): return len(self.jobs)
    @property
    def n_ops_total(self): return sum(len(j.ops) for j in self.jobs)

    # ---- 染色体编码 ----
    # 基于常用编码：两个基因向量
    # 1) op_seq: 长度 = 总工序数，由作业id组成（每作业重复出现其工序数），决定工序出现顺序
    # 2) mac_seq: 长度 = 总工序数，对应每一次出现的“工序”，选择一个可行机器索引（在该工序的可行列表中的下标）
    # 解码时：遇到作业 j 的第 t 次出现 => 安排该作业的第 t 道工序，机器从 mac_seq 对应位置选。

    def feasible_machines(self, j, k) -> List[int]:
        return list(self.jobs[j].ops[k].machine_time.keys())

    def decode(self, op_seq: List[int], mac_seq: List[int]) -> Tuple[float, Dict[str,List[Tuple[float,float,str]]]]:
        # 状态
        job_next = np.zeros(self.n_jobs, dtype=int)           # 下一道待执行的工序号
        job_ready = np.zeros(self.n_jobs, dtype=float)        # 作业就绪时间（上道完成）
        mach_ready = np.zeros(self.m, dtype=float)            # 机器就绪时间
        gantt = {f"M{mm}":[] for mm in range(self.m)}

        mac_ptr = 0
        for j in op_seq:
            k = job_next[j]                                   # 该作业的第k道工序
            if k >= len(self.jobs[j].ops):
                continue  # 已完成（保护性）
            feas = self.feasible_machines(j, k)               # 可行机列表
            # 从 mac_seq 给的指示取一个（按索引映射到可行机 id）
            idx = mac_seq[mac_ptr] % len(feas)
            mm = feas[idx]
            pt = float(self.jobs[j].ops[k].machine_time[mm])
            st = max(job_ready[j], mach_ready[mm])
            en = st + pt
            job_ready[j] = en; mach_ready[mm] = en
            job_next[j] += 1
            gantt[f"M{mm}"].append((st, pt, f"J{j}-O{k}"))
            mac_ptr += 1

        Cmax = float(np.max(job_ready))
        return Cmax, gantt

# ---- GA for FJSP ----

class GA_FJSP:
    def __init__(self, problem:FJSP, pop=80, gens=400, pc=0.9, pm=0.2, seed=0):
        self.pb=problem; self.pop=pop; self.gens=gens; self.pc=pc; self.pm=pm; seed_all(seed)
        # 预生成 op_seq 基础串（作业id按工序数重复）
        self.base_seq=[]
        for j, job in enumerate(self.pb.jobs):
            self.base_seq += [j]*len(job.ops)

    def rand_individual(self):
        # op_seq: shuffle base_seq
        op_seq = self.base_seq[:]; random.shuffle(op_seq)
        # mac_seq: 每次出现对应一个工序，先随机在可行机列表内选一个下标（统一存“下标”）
        mac_seq=[]
        times_seen = [0]*self.pb.n_jobs
        for j in op_seq:
            k = times_seen[j]
            feas = self.pb.feasible_machines(j, k)
            mac_seq.append(random.randrange(len(feas)))
            times_seen[j]+=1
        return (op_seq, mac_seq)

    def crossover_ppx(self, a_seq, b_seq):
        """Precedence Preservative Crossover for repeated job-ids序列：PPX-like
           随机抽取位置次序，逐步从父代选择下一个合法作业id（保持每作业出现次数不变）
        """
        n = len(a_seq)
        # 统计每作业总计需要次数
        need = {}
        for x in a_seq:
            need[x]=need.get(x,0)+1
        left = dict(need)
        res=[]
        a_ptrs = {j:[] for j in set(a_seq)}
        b_ptrs = {j:[] for j in set(a_seq)}
        for i,x in enumerate(a_seq): a_ptrs[x].append(i)
        for i,x in enumerate(b_seq): b_ptrs[x].append(i)
        a_pos, b_pos = 0, 0
        while len(res)<n:
            if random.random()<0.5:
                # 从a选择下一个在need中未用尽的作业
                while a_pos<n and left.get(a_seq[a_pos],0)==0: a_pos+=1
                if a_pos<n:
                    j = a_seq[a_pos]; res.append(j); left[j]-=1; a_pos+=1
                else:
                    # fallback: 从b取
                    while b_pos<n and left.get(b_seq[b_pos],0)==0: b_pos+=1
                    j = b_seq[b_pos]; res.append(j); left[j]-=1; b_pos+=1
            else:
                while b_pos<n and left.get(b_seq[b_pos],0)==0: b_pos+=1
                if b_pos<n:
                    j = b_seq[b_pos]; res.append(j); left[j]-=1; b_pos+=1
                else:
                    while a_pos<n and left.get(a_seq[a_pos],0)==0: a_pos+=1
                    j = a_seq[a_pos]; res.append(j); left[j]-=1; a_pos+=1
        return res

    def crossover_uniform_mac(self, a_mac, b_mac):
        # 均匀交叉机器下标向量
        return [ (a if random.random()<0.5 else b) for a,b in zip(a_mac, b_mac) ]

    def mutate(self, op_seq, mac_seq):
        # 序列小范围交换
        if random.random()<0.5:
            i,j = random.sample(range(len(op_seq)),2)
            op_seq[i],op_seq[j] = op_seq[j],op_seq[i]
        # 机器分配小范围变异：随机挑若干位置，重采样可行机下标
        if random.random()<0.7:
            times_seen=[0]*self.pb.n_jobs
            for i,j in enumerate(op_seq):
                k = times_seen[j]
                if random.random()<0.2:
                    feas = self.pb.feasible_machines(j, k)
                    mac_seq[i] = random.randrange(len(feas))
                times_seen[j]+=1

    def run(self, verbose=True):
        pop=[self.rand_individual() for _ in range(self.pop)]
        fits=[self.pb.decode(ind[0], ind[1])[0] for ind in pop]
        best_hist=[]
        for g in range(1, self.gens+1):
            new=[]
            while len(new) < self.pop:
                i,j = random.randrange(self.pop), random.randrange(self.pop)
                p1 = pop[i] if fits[i]<fits[j] else pop[j]
                i,j = random.randrange(self.pop), random.randrange(self.pop)
                p2 = pop[i] if fits[i]<fits[j] else pop[j]
                c1=(p1[0][:], p1[1][:]); c2=(p2[0][:], p2[1][:])
                if random.random()<self.pc:
                    c1 = (self.crossover_ppx(p1[0], p2[0]),
                          self.crossover_uniform_mac(p1[1], p2[1]))
                    c2 = (self.crossover_ppx(p2[0], p1[0]),
                          self.crossover_uniform_mac(p2[1], p1[1]))
                if random.random()<self.pm: self.mutate(*c1)
                if random.random()<self.pm: self.mutate(*c2)
                new += [c1,c2]
            pop = new[:self.pop]
            fits=[self.pb.decode(ind[0], ind[1])[0] for ind in pop]
            best=min(fits); best_hist.append(best)
            if verbose and g%50==0:
                print(f"[GA_FJSP] Gen {g:4d} Best Cmax={best:.2f}")
        idx=int(np.argmin(fits))
        return pop[idx], fits[idx], best_hist

# ============ Demo 构造随机实例 ============

def build_random_fsp(n_jobs=20, n_machines=5, pt_low=1, pt_high=10, seed=0) -> FSP:
    rng = np.random.RandomState(seed)
    p = rng.randint(pt_low, pt_high+1, size=(n_jobs, n_machines)).astype(float)
    return FSP(p)

def build_random_hfs(n_jobs=20, stages=3, Ms=[2,2,2], pt_low=1, pt_high=10, seed=0) -> HFS:
    rng = np.random.RandomState(seed)
    proc=[]
    for s in range(stages):
        proc.append(rng.randint(pt_low, pt_high+1, size=(n_jobs,)).astype(float))
    return HFS(proc_time=proc, Ms=Ms)

def build_random_fjsp(n_jobs=10, m=6, ops_per_job=4, k_machines_per_op=2, pt_low=1, pt_high=9, seed=0) -> FJSP:
    rng = np.random.RandomState(seed)
    jobs=[]
    for _ in range(n_jobs):
        ops=[]
        for _o in range(ops_per_job):
            feas = sorted(rng.choice(range(m), size=k_machines_per_op, replace=False))
            mt = {mm:int(rng.randint(pt_low, pt_high+1)) for mm in feas}
            ops.append(Operation(machine_time=mt))
        jobs.append(JobF(ops=ops))
    return FJSP(jobs=jobs, m=m)

# ============ main: 运行不同算法的示例 ============

def demo_fsp():
    print("=== FSP: GA & SA 示例 ===")
    fsp = build_random_fsp(n_jobs=20, n_machines=5, seed=42)
    # GA
    ga = GA_FSP(fsp, pop=60, gens=300, seed=1)
    best_perm, best_val, hist = ga.run(verbose=True)
    print(f"[GA_FSP] Best makespan = {best_val:.2f}")
    plot_convergence(hist, title="FSP - GA Convergence")
    gantt = fsp.decode_gantt(best_perm)
    plot_gantt(gantt, title=f"FSP GA Gantt Cmax={best_val:.1f}")
    # SA
    sa = SA_FSP(fsp, iters=4000, seed=2)
    best_perm2, best_val2, hist2 = sa.run(verbose=True)
    print(f"[SA_FSP] Best makespan = {best_val2:.2f}")
    plot_convergence(hist2, title="FSP - SA Convergence")
    gantt2 = fsp.decode_gantt(best_perm2)
    plot_gantt(gantt2, title=f"FSP SA Gantt Cmax={best_val2:.1f}")

def demo_hfs():
    print("=== HFS: GA 示例 ===")
    hfs = build_random_hfs(n_jobs=24, stages=3, Ms=[2,3,2], seed=7)
    ga = GA_HFS(hfs, pop=70, gens=350, seed=3)
    (perm,assign), best_val, hist = ga.run(verbose=True)
    print(f"[GA_HFS] Best makespan = {best_val:.2f}")
    plot_convergence(hist, title="HFS - GA Convergence")
    Cmax, gantt = hfs.decode(perm, assign)
    plot_gantt(gantt, title=f"HFS GA Gantt Cmax={Cmax:.1f}")

def demo_fjsp():
    print("=== FJSP: GA 示例 ===")
    fjsp = build_random_fjsp(n_jobs=12, m=6, ops_per_job=4, k_machines_per_op=2, seed=11)
    ga = GA_FJSP(fjsp, pop=90, gens=400, seed=5)
    (op_seq, mac_seq), best_val, hist = ga.run(verbose=True)
    print(f"[GA_FJSP] Best makespan = {best_val:.2f}")
    plot_convergence(hist, title="FJSP - GA Convergence")
    Cmax, gantt = fjsp.decode(op_seq, mac_seq)
    plot_gantt(gantt, title=f"FJSP GA Gantt Cmax={Cmax:.1f}")

if __name__ == "__main__":
    seed_all(0)
    # 任选一个或多个 demo 运行
    demo_fsp()
    demo_hfs()
    demo_fjsp()
